# @Time   : 2023/3/29 11:48
# @Author : Fioman
# @Phone  : 13149920693
# @Tips   : Talk is Cheap,Show me the code! ^_^^_^

import datetime
import os
from settings import OCR_GET_RES_DIR
from tools.config_params import ConfigParams
import cv2 as cv
import numpy as np
from cnocr import CnOcr
import time

"""
OCR操作相关的函数
"""
cp = ConfigParams()

# 使用模型介绍
detectionModelName = "en_PP-OCRv3_det"  # 检测模型

# detectionModelName = "ch_PP-OCRv3_det"  # 检测模型
recModelName = "en_PP-OCRv3"  # 识别模型
# recModelName = "densenet_lite_136-fc"  # 识别模型

candAlphabet = "0123456789-"  # 要识别的候选集合
# resizeShapeList = [(768, 768), (512, 768), (768, 1024), (1024, 1024)]  # 检测前图像的resize,(height,width)
resizeShapeList = [(256, 512), (512, 768), (132, 768), (768, 1024)]  # 检测钱图像的resize,(height,width)
minBoxSize = 5  # 如果检测出的文本框高度或者宽度低于此值,文本框会被过滤掉.
boxScoreThresh = 0.3  # 过滤掉低于此分的结果

textLenPrint = 12  # 激光喷码的字符个数
textWidthPix = 11.4  # 单个字符的宽度为11.4,如果和这个误差最小的字符算是有效字符.
textHeightPix = 100  # 字符高度像素


def image_show(name, image, showState=False):
    """
    :param name:
    :param image:
    :param showState:
    :return:
    """
    if showState:
        sizeScaleWidth = 0.2
        sizeScaleHeight = 0.2
        newWidth = int(image.shape[1] / sizeScaleWidth)
        newHeight = int(image.shape[0] / sizeScaleHeight)
        showImage = cv.resize(image, (newWidth, newHeight), cv.INTER_AREA)
        cv.namedWindow("{}".format(name))
        cv.imshow("{}".format(name), showImage)
        cv.waitKey(0)


def get_year_month_day_from_dateStr(dataStr):
    """
    根据形如这样的字符串 20230328
    :param dataStr:
    :return:
    """
    resList = []
    try:
        resList.append(dataStr[:4])
        resList.append(dataStr[4:6])
        resList.append(dataStr[6:8])
    except Exception as e:
        resList = []
    finally:
        return resList


def get_datetime_now_str():
    """
    获取当前日期的字符串
    :return:
    """
    datetimeNow = datetime.datetime.now()
    return datetimeNow.strftime("%Y%m%d")


def day_add_and_sub(oldData):
    """
    这个函数主要的功能是获取当前日期的前后日期,防止日期有问题,然后板子找不到的情况.
    :param oldData:
    :return:
    """
    try:
        # 1. 先获取oldData中的年月日
        dateList = get_year_month_day_from_dateStr(oldData)
        # 1.1 日期格式不对直接就按原来的日期返回
        if not dateList:
            return oldData
        # 1.2 解压出来年月日
        year, month, day = dateList
        currentDate = datetime.datetime(int(year), int(month), int(day))
        nextDate = currentDate + datetime.timedelta(days=+1)
        prevDate = currentDate + datetime.timedelta(days=-1)
        nextDateStr = nextDate.strftime("%Y%m%d")
        prevDateStr = prevDate.strftime("%Y%m%d")
        return [prevDateStr]
    except Exception as e:
        debugInfo = "in day_add_and_sub() Error: {}".format(str(e))
        print(debugInfo)
        return []


def OCR_get_ng_or_ok_res_new(dateStr, boardIndex):
    """
    快速定位文件,如果不存在就换前一天的目录.
    :param dateStr:
    :param boardIndex:
    :return:
    """
    try:
        filePath = os.path.join(OCR_GET_RES_DIR, dateStr)
        if os.path.exists(filePath):
            # true文件的大小,先判断文件是否存在
            truFileList = ["{}_b_True".format(boardIndex), "{}_t_True".format(boardIndex)]
            trueFilePath = [os.path.join(filePath, truFileList[0]), os.path.join(filePath, truFileList[1])]
            trueFileSizeList = []
            for filePathTemp in trueFilePath:
                if os.path.exists(filePathTemp):
                    trueFileSizeList.append(os.path.getsize(filePathTemp))
                else:
                    trueFileSizeList.append(-1)
            if trueFileSizeList[0] == -1 or trueFileSizeList[1] == -1:
                if trueFileSizeList[0] == -1:
                    return "NG", "ID({})True文件,人工确认OK的文件({})找不到,NG不通过!".format(boardIndex, truFileList[0])
                else:
                    return "NG", "ID({})True文件,人工确认OK的文件({})找不到,NG不通过!".format(boardIndex, truFileList[1])
            if trueFileSizeList[0] > 8:
                return "NG", "ID({}) 对应的True文件,人工确认NG,判断依据: {}文件内容大于8".format(boardIndex, truFileList[0])
            if trueFileSizeList[1] > 8:
                return "NG", "ID({}) 对应的True文件,人工确认NG,判断依据: {}文件内容大于8".format(boardIndex, truFileList[1])
            return "ok", "ID({})对应的True文件_{}字节等于8,人工确认OK".format(boardIndex, truFileList)
        else:
            if dateStr != get_datetime_now_str():
                dateUsed = get_datetime_now_str()
            else:
                dateUsed = day_add_and_sub(dateStr)[0]
            filePath = os.path.join(OCR_GET_RES_DIR, dateUsed)
            if os.path.exists(filePath):
                # true文件的大小,先判断文件是否存在
                truFileList = ["{}_b_True".format(boardIndex), "{}_t_True".format(boardIndex)]
                trueFilePath = [os.path.join(filePath, truFileList[0]), os.path.join(filePath, truFileList[1])]
                trueFileSizeList = []
                for filePathTemp in trueFilePath:
                    if os.path.exists(filePathTemp):
                        trueFileSizeList.append(os.path.getsize(filePathTemp))
                    else:
                        trueFileSizeList.append(-1)
                if trueFileSizeList[0] == -1 or trueFileSizeList[1] == -1:
                    if trueFileSizeList[0] == -1:
                        return "NG", "ID({})True文件,人工确认OK的文件({})找不到,NG不通过!".format(boardIndex, truFileList[0])
                    else:
                        return "NG", "ID({})True文件,人工确认OK的文件({})找不到,NG不通过!".format(boardIndex, truFileList[1])
                if trueFileSizeList[0] > 8:
                    return "NG", "ID({}) 对应的True文件,人工确认NG,判断依据: {}文件内容大于8".format(boardIndex, truFileList[0])
                if trueFileSizeList[1] > 8:
                    return "NG", "ID({}) 对应的True文件,人工确认NG,判断依据: {}文件内容大于8".format(boardIndex, truFileList[1])
                return "ok", "ID({})对应的True文件_{}字节等于8,人工确认OK".format(boardIndex, truFileList)
            return "PathError", "目录{}不存在".format(filePath)
    except Exception as e:
        debugInfo = "ID({}) in OCR_get_ng_or_ok_res() error: {}".format(boardIndex, (e))
        return "ok", debugInfo


def OCR_get_ng_or_ok_res_old(dateStr, boardIndex):
    """
    根据目录和板编号获取对应的板子的信息,如果没有True文件,默认是OK的,
    如果找不到对应的目录就找相邻的天的目录
    :param dateStr: 日期字符串,子文件夹,一般是当天的日期. 20230304这种类型的
    :param boardIndex:
    :return:
    """
    try:
        filePath = os.path.join(OCR_GET_RES_DIR, dateStr)
        if os.path.exists(filePath):
            subDirList = [file for file in os.listdir(filePath) if os.path.isdir(os.path.join(filePath, file))]
        else:
            subDirList = []
        # random.shuffle(subDirList)
        # print("打乱排序之前的列表: {}".format(subDirList))
        # subDirList.sort()
        # print("打乱排序之后的列表: {}".format(subDirList))
        # 如果板编号不在当前的日期里面,就去前面一天和后一天去找
        print(subDirList)
        if boardIndex not in subDirList:
            nextDate, prevDate = day_add_and_sub(dateStr)
            filePathNext = os.path.join(OCR_GET_RES_DIR, nextDate)
            print(filePathNext)
            filePath = filePathNext
            if os.path.exists(filePath):
                subDirList = [file for file in os.listdir(filePathNext) if os.path.isdir(os.path.join(filePathNext, file))]
            else:
                subDirList = []
            if boardIndex not in subDirList:
                filePathPrev = os.path.join(OCR_GET_RES_DIR, prevDate)
                filePath = filePathPrev
                subDirList = [file for file in os.listdir(filePath) if os.path.isdir(os.path.join(filePath, file))]
                if boardIndex not in subDirList:
                    return "ok", "(ID)文件目录找不多,查找的目录为{}".format(boardIndex, [dateStr, nextDate, prevDate])
        # 寻找的文件名称
        trueFileList = ["{}_b_True".format(boardIndex), "{}_t_True".format(boardIndex)]
        # 确定好filePath之后就是寻找对应的文件了.
        trueFileFindedList = [file for file in os.listdir(filePath) if file in trueFileList]
        print("TrueFileList: {}".format(trueFileFindedList))
        if len(trueFileFindedList) < 1:
            return "ok", "ID({})True文件,人工确认OK和NG的文件目录找不到,默认是OK通过!".format(boardIndex)

        # 如果找到了true文件,就遍历这个true文件,如果发现了字节个数大于8个就NG.
        for file in trueFileFindedList:
            if os.path.getsize(os.path.join(filePath, file)) > 8:
                return "NG", "ID({}) 对应的True文件,人工确认NG,判断依据: {}文件内容大于8".format(boardIndex, file)
        return "ok", "ID({})对应的True文件_{}字节等于8,人工确认OK".format(boardIndex, trueFileFindedList)

    except Exception as e:
        debugInfo = "ID({}) in OCR_get_ng_or_ok_res() error: {}".format(boardIndex, (e))
        return "ok", debugInfo


def get_image_used_for_ocr(image, xSplitStart=0.5, xSplitEnd=1):
    """
    获取彩色图,先截取X轴,然后是逆时针旋转
    :param image
    :param xSplit:
    :return:
    """
    try:
      print(int(image.shape[0]), int(image.shape[0]))
      imgColorUsed = cv.rotate(image, cv.ROTATE_90_COUNTERCLOCKWISE)
      imgColorUsed = imgColorUsed[int(imgColorUsed.shape[0] * xSplitStart):int(imgColorUsed.shape[0]), :]
      print(int(imgColorUsed.shape[0]), int(imgColorUsed.shape[1]))
    except Exception as e:
      print("出现异常: {}".format(str(e)))
      imgColorUsed = image.copy()
    # 测试的时候先按照原图进行返回
    return imgColorUsed


def get_text_list(dictRes):
    if len(dictRes) == 0:
        return []

    resList = []
    textResList = []
    for textInfo in dictRes:
        textFinded = textInfo.get("text")
        pointList = textInfo.get("position")
        widthFindedPix = abs(pointList[0][0] - pointList[1][0])
        heightFindedPix = abs(pointList[0][1] - pointList[2][1])
        textCenterX = (pointList[0][0] + pointList[1][0]) / 2
        textLen = len(textFinded)
        # 检测字符串是否有效
        if textLen == 12:
            textRes = textFinded
            widthError = abs(widthFindedPix / len(textRes) - textWidthPix)
            temp = (textRes, textLen, widthFindedPix, heightFindedPix, widthError)
            resList.append(temp)
        elif 4 <= textLen <= 5:
            temp = (textFinded, textLen, widthFindedPix, heightFindedPix, textCenterX)
            textResList.append(temp)
            # print("textFinded: {}, perTextWidth: {}".format(textFinded,widthFindedPix / len(textRes)))
            # if 8 <= len(textRes) <= 9:
            #     temp = (textRes,len(textRes),widthFindedPix,heightFindedPix)
            #     resList.append(temp)
            #     return resList
    # 筛选出来合适的字体
    firstPart = []
    secondPart = []
    distanceMax = 90
    if len(textResList) >= 2:
        textResList.sort(key=lambda x: x[4])
        for index, text in enumerate(textResList):
            if text != textResList[-1]:
                nextText = textResList[index + 1]
                distanceError = abs(nextText[-1] - text[-1])
                if distanceError < distanceMax:
                    if firstPart == []:
                        firstPart = text
                        secondPart = nextText
                        break
    if firstPart != [] and secondPart != []:
        widthFindedPix = firstPart[2] + secondPart[2]
        textRes = firstPart[0] + secondPart[0]
        widthError = abs(widthFindedPix / len(textRes) - textWidthPix)
        temp = (
        firstPart[0] + secondPart[0], len(firstPart[0]) + len(secondPart[0]), firstPart[1], firstPart[2], widthError)
        if temp[1] == 12:
            resList.append(temp)
    resList.sort(key=lambda x: x[4])
    return resList


def check_text_is_ok(text):
    """
    检测这个text是否满足要求
    :param text:
    :return:
    """
    textVal, textLen, widthPix, heightPix, widthPixError = text
    if textLen == 12:
        if textVal[0] in "235689":
            textValList = list(textVal)
            textValList[0] = "0"
            textVal = "".join(textValList)
        elif textVal[0] in "47":
            textValList = list(textVal)
            textValList[0] = "1"
            textVal = "".join(textValList)

        boardIndex = int(textVal[:-7])
        dateStr = textVal[-4:]
        return "ok", [textVal, boardIndex, dateStr]
    elif textLen == 100:
        if textVal[-4:] == get_datetime_now_str()[-4:]:
            textVal = "0" + textVal
            boardIndex = int(textVal[:-4])
            dateStr = textVal[-4:]
        else:
            textVal = textVal + get_datetime_now_str()[-1]
            boardIndex = int(textVal[:-4])
            dateStr = textVal[-4:]
        return "okByFillZero", [textVal, boardIndex, dateStr]
    return "error", ["识别失败", None, None]


def find_text_by_cnocr(image):
    """
    识别图像
    :param image:
    :return:
    """
    # imageUsed = get_image_used_for_ocr(image)
    imageUsed = image.copy()
    # imageUsed = cv.cvtColor(imageUsed,cv.COLOR_BGR2GRAY)
    # image_show("ImageUsed", imageUsed)
    # 做一些图像预处理,比如高斯模糊,还有开操作.
    # gaussianRes =cv.blur(imageUsed,(3,3))
    # erode = cv.erode(imageUsed,(3,3),iterations=3)
    # image_show("OpennedRes",erode)
    # hist = cv.equalizeHist(imageUsed)
    # image_show("Hist",hist)

    ocr = CnOcr(rec_model_name=recModelName,
                det_model_name=detectionModelName,
                cand_alphabet=candAlphabet)
    resFindedList = []
    for resizeShape in resizeShapeList:
        dictRes = ocr.ocr(imageUsed, resized_shape=resizeShape)
        for res in dictRes:
            if res["text"] != "":
                print("DictRes, text: {}, score:{}".format(res["text"], res["score"]))
        # print("dictRes: {}".format(dictRes))
        textList = get_text_list(dictRes)  # 通过dictRes获取字符信息
        print("使用resize:{},识别有效结果: {}".format(resizeShape, textList))
        if len(textList) >= 1:
            for text in textList:
                res, textRes = check_text_is_ok(text)
                print("字符有效检测结果: res: {},textRes: {}".format(res, textRes))
                if res == "ok":
                    return "ok", "识别12个字符成功", textRes
                elif res == "okByFillZero":
                    resFindedList.append(textRes)

    if len(resFindedList) >= 1:
        return "okByFillZero", "识别到8个字符,前面字符填充为0", resFindedList[0]
    return "error", "识别错误,没有找到对应字符", []


if __name__ == '__main__':
    #
    # nowStr = get_datetime_now_str()
    # print("nowStr: " + nowStr)
    # [nowStr].extend(day_add_and_sub(nowStr))
    # print("nosStr: " + nowStr)
    #
    # dateDir  = "20230610"
    # try:
    #
    #     try:
    #         monthError = int(dateDir[-4:-2]) - int(nowStr[-4:-2])
    #         dateError = int(dateDir[-2:]) - int(nowStr[-2:])
    #         # 计算的日期,比当前的日期还要超前,理论上是不存在的.
    #         if (monthError == 0 and dateError > 1) or (abs(monthError) >= 2) or(abs(dateError) > 5):
    #             dateDir = nowStr
    #             print("日期计算错误,采用当前日期. monthError: {},dateError: {}".format(monthError,dateError))
    #     except Exception as e:
    #        print("日期计算错误,继续往下执行")
    # except:
    #     print("异常会影响我吗")
    #
    #
    # while True:
    #     pass

    '''
    ********** 这部分是测试文件目录读取功能 ************
    '''

    # dateStr = "20230407"
    # indexList = ["5649"]
    # timeReadStart = time.time()
    # for index in indexList:
    #     res, resInfo = OCR_get_ng_or_ok_res_new(dateStr, index)
    #     print("res: {}, {}".format(res, resInfo))
    # timeReadEnd = time.time()
    #
    # print("readFileTimeUsed: {}".format(timeReadEnd - timeReadStart))

    '''
        ********** 这部分是测试文件目录读取功能 ************
    '''

    # 下面这部分是测试算法功能

    if os.path.exists("res.txt"):
        os.remove("res.txt")

    with open("res.txt", "w", encoding="utf-8") as f:
        f.write("识别结果   字符    板号   日期  文件名称\n")

    filePath = r"D:\demarcate_color_res"
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "ok")
    keepOkResPath = os.path.join(keepDir, "ok_res")
    keepFailPath = os.path.join(keepDir, "fail")

    dirList = [keepOkPath, keepOkResPath, keepFailPath]
    for dirDemo in dirList:
        if not os.path.exists(dirDemo):
            os.mkdir(dirDemo)

    fileNames = []
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:
                continue
            for file in files:
                if file[-4:] != ".bmp":
                    continue
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        print("*" * 10 + "第 {} 张图 ({})".format(index + 1, fileName) + "*" * 10)
        calcStart = time.time()
        filePathReal = os.path.join(filePath, fileName)
        imageOriginal = cv.imread(filePathReal, cv.IMREAD_COLOR)
        # imageGray = cv.cvtColor(imageOriginal,cv.COLOR_BGR2GRAY)
        res, info, textResList = find_text_by_cnocr(imageOriginal)

        if res == "ok" or res == "okByFillZero":
            with open("res.txt", "a") as f:
                f.write("{}, {}, {}, {}, {}\n".format(res, textResList[0], textResList[1], textResList[2], fileName))
            if res == "ok":
                print("识别成功: 识别结果: {}".format(textResList))
            else:
                print("识别成功: 通过前面填充0,识别结果: {}".format(textResList))
            os.remove(filePathReal)
            imageRes = cv.putText(imageOriginal.copy(), textResList[0], (50, 30), cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            okResPathReal = os.path.join(keepOkResPath, fileName)
            if os.path.exists(okResPathReal):
                os.remove(okResPathReal)
            cv.imwrite(okResPathReal, imageRes)
            cv.imwrite(os.path.join(keepOkPath, fileName), imageOriginal)
            calcEnd = time.time()
            print("算法耗时: {:3f}".format(calcEnd - calcStart))
            continue
        else:
            with open("res.txt", "a") as f:
                f.write("{} {} {},{}, {}\n".format(res, 0, 0, 0, fileName))
            print("识别失败xxxx")
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepFailPath, fileName), imageOriginal)
            calcEnd = time.time()
            print("算法耗时: {:3f}".format(calcEnd - calcStart))
